package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"runtime/pprof"
	_ "runtime/pprof"
	"sync"
	"sync/atomic"
	"time"

	"github.com/nsqio/go-nsq"
)

const (
	topicPrefix = "sub_bench_"
)

var (
	runfor     = flag.Duration("runfor", 10*time.Second, "duration of time to run")
	lookupAddr = flag.String("nsqlookupd-http-address", "127.0.0.1:4161", "<addr>:<port> to connect to nsqlookupd http")
	size       = flag.Int("size", 200, "size of messages")
	topicCount = flag.Int("topic-count", 1, "topic count to receive messages on")
	channel    = flag.String("channel", "ch", "channel to receive messages on")
	deadline   = flag.String("deadline", "", "deadline to start the benchmark run")
	rdy        = flag.Int("rdy", 2500, "RDY count to use")
	cpuProfile = flag.String("cpu-profile", "", "write cpu profile to file")
)

var totalMsgCount int64

func main() {
	flag.Parse()
	var wg sync.WaitGroup

	if *cpuProfile != "" {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	log.SetPrefix("[bench_reader] ")

	goChan := make(chan int)
	rdyChan := make(chan int)
	workers := runtime.GOMAXPROCS(0)
	log.Printf("workers: %d", workers)
	for j := 0; j < workers; j++ {
		wg.Add(1)
		go func(id int) {
			subWorker(*runfor, workers, *lookupAddr, *topicCount, *channel, rdyChan, goChan, id)
			wg.Done()
		}(j)
		<-rdyChan
	}

	if *deadline != "" {
		t, err := time.Parse("2006-01-02 15:04:05", *deadline)
		if err != nil {
			log.Fatal(err)
		}
		d := t.Sub(time.Now())
		log.Printf("sleeping until %s (%s)", t, d)
		time.Sleep(d)
	}

	start := time.Now()
	close(goChan)
	wg.Wait()
	end := time.Now()
	duration := end.Sub(start)
	tmc := atomic.LoadInt64(&totalMsgCount)
	mbps := float64(tmc*int64(*size)) / duration.Seconds() / 1024 / 1024
	opspms := float64(tmc) / duration.Seconds() / 1000
	log.Printf("duration: %s - %.01fmb/s - %.01fops/ms - %.03fus/op",
		duration, mbps, opspms, float64(duration/time.Microsecond)/float64(tmc))
	log.Printf("Result: %d\t%.01f\t%.01f", *topicCount, mbps, opspms)
}

type myMessageHandler struct {
	msgCount int64
}

// HandleMessage implements the Handler interface.
func (h *myMessageHandler) HandleMessage(m *nsq.Message) error {
	h.msgCount++
	// Returning nil will automatically send a FIN command to NSQ to mark the message as processed.
	return nil
}

func subWorker(td time.Duration, workers int, lookupAddr string, topicCount int, channel string, rdyChan chan int, goChan chan int, id int) {
	// Instantiate a consumer that will subscribe to the provided channel.
	config := nsq.NewConfig()
	config.DialTimeout = time.Second
	config.MaxInFlight = 1000

	handlers := make([]myMessageHandler, topicCount)
	consumers := make([]*nsq.Consumer, topicCount)
	for i := 0; i < topicCount; i++ {
		topic := fmt.Sprintf("%s%d", topicPrefix, i+1)
		consumer, err := nsq.NewConsumer(topic, channel, config)
		if err != nil {
			panic(err)
		}
		consumer.SetLogger(nil, 0)

		consumer.AddHandler(&handlers[i])
		consumers[i] = consumer
	}

	rdyChan <- 1
	<-goChan

	var wgGo sync.WaitGroup
	for i := 0; i < topicCount; i++ {
		wgGo.Add(1)
		consumer := consumers[i]
		go func() {
			err := consumer.ConnectToNSQLookupd(lookupAddr)
			if err != nil {
				panic(err)
			}
			wgGo.Done()
		}()
	}

	time.Sleep(td)
	wgGo.Wait()

	var wgStop sync.WaitGroup
	for i := 0; i < topicCount; i++ {
		wgStop.Add(1)
		// Gracefully stop the consumer.
		consumer := consumers[i]
		go func() {
			consumer.Stop()
			wgStop.Done()
		}()
	}
	wgStop.Wait()
	for i := 0; i < topicCount; i++ {
		atomic.AddInt64(&totalMsgCount, handlers[i].msgCount)
	}
}
